from nested_gen import *
import random

def flipCol(mat, index, n, m, p):
  """OVERVIEW: Finds the position of the isocline within the given
               column; randomly finds a cell above and a cell below
               whose values are NOT unexpected; and flips their values.
               If either side of the isocline contains only unexpected
               values, no flip is performed (because a flip would lower
               the matrix temperature).
  PARAMETERS:
              mat, a site-species presence matrix
              index, the index of the column whose cells to flip
              n, the number of sites
              m, the number of species
              p, the isocline parameter

  RETURNS:  If a flip has been performed, returns the indices of the
            affected rows.
            Otherwise, returns -1.
            (Changes to the matrix are made directly)
  """
  # To prevent flips that lower matrix temperature,
  # get positions of cells that are NOT unexpected
  cells = unflipped(mat, index, n, m, p, True)

  if(len(cells[0]) > 0 and len(cells[1]) > 0):
    toflip_above = random.sample(cells[0],1)
    toflip_below = random.sample(cells[1],1)
  else: # everything on at least one side of the isocline has been flipped
    return -1

  # transpose mat so that indices run along columns, not rows
  mat = mat.transpose()

  # flip cells
  mat[index, toflip_above] = 0
  mat[index, toflip_below] = 1

  # re-transpose to get original configuration
  mat.transpose()
  # return affected rows
  return(toflip_above, toflip_below)

# end flipCol


def flipRow(mat, index, n, m, p):
  """OVERVIEW: Finds the position of the isocline within the given row;
               randomly finds a cell to the left and a cell to the
               right whose values are NOT unexpected; and flips their
               values.
               If either side of the isocline contains only unexpected
               values, no flip is performed (because a flip would lower
               the matrix temperature).
  PARAMETERS:
               mat, a site-species presence matrix
               index, the index of the row whose cells to flip
               n, the number of sites
               m, the number of species
               p, the isocline parameter

  RETURNS: If a flip has been performed, returns nothing.
           Otherwise, returns -1.
           (Changes to the matrix are made directly)
  """
  # get positions of cells that haven't been flipped yet
  cells = unflipped(mat, index, n, m, p, False)

  if(len(cells[0]) > 0 and len(cells[1]) > 0):
    toflip_left  = random.sample(cells[0],1)
    toflip_right = random.sample(cells[1],1)
  else: # everything on at least one side of the isocline has been flipped
    return -1

  # flip cells
  mat[index, toflip_left]  = 0
  mat[index, toflip_right] = 1

  return()

# end flipRow


def unflipped(mat, index, n, m, p, col = False):
  """OVERVIEW: Finds all cells in the given row or column that are NOT
               unexpected.
 PARAMETERS:
            mat, a site-species presence matrix
            index, the index of the row/col whose cells to flip
            n, the number of sites
            m, the number of species
            p, the isocline parameter
            col, a boolean indicating whether to flip along columns
 RETURNS:
         A tuple containing two lists of positions of cells that
         are NOT unexpected.
         The [0] index contains cell positions above or left of the isocline.
         The [1] index contains cell positions below or right of the isocline.
  """
  # determine whether to look at row or column
  if(col == True):
    # transpose so that indices run along columns
    mat = mat.transpose()

    # get total number of cells above the isocline (i.e. to the left
    # after the transpose).
    cells_to_left = n - round(isocline((index+.5)/m,p,n,m)*n)

  else: # look across the row for total cells to left
    cells_to_left = m - round(isocline((index+.5)/n,p,m,n)*m)

  # pull out the desired row/col
  row = mat[index]

  # get unflipped (nonzero) cells to the left
  left_unflipped = numpy.nonzero(row[0:cells_to_left])[0]

  # get unflipped (zero) cells to the right
  cells_to_right = len(row) - cells_to_left
  right_unflipped = numpy.nonzero(row[cells_to_left:]-1)[0]
  # In right_unflipped, indexing starts at isocline.  Add
  # cells_to_left to indices to get the correct values
  right_unflipped = right_unflipped + cells_to_left

  # return list of indices of unflipped cells in a tuple
  return((left_unflipped, right_unflipped))

# end unflipped


def totUnexp(mat,n,m,p):
  """OVERVIEW: Calculates the total unexpectedness of a matrix.
  PARAMETERS:
             mat, a site-species presence matrix
             n, the number of sites
             m, the number of species
             p, the isocline parameter
  RETURNS:
         Matrix unexpectedness (a float)
  """
  # sum the unexpectedness in the cells of a matrix
  unexp = 0 # initial unexpectedness
  row_indices = set(range(m)) # the set of indices within each row
  for i in range(n):
    # Get indices of unflipped cells in row i
    unflipped_i = unflipped(mat, i, n, m, p, col = False)

    # now get indices of cells that have been flipped
    unflipped_set = set(numpy.append(unflipped_i[0], unflipped_i[1]))
    flipped_i = row_indices.difference(unflipped_set)

    # calculate the unexpectedness of flipped cells in row i
    for f in flipped_i:
      unexp = unexp + unexpected(n,m,p,i,f) # add to total unexpectedness
    # end for
  # end for
  return (1./(n*m))*unexp # scale result by matrix size
# end totUnexp()


def hotMess(n,m,phi, T, err = 1, max_iter = 100):
  """OVERVIEW: Creates a site-species presence matrix with a given fill
               and temperature.
  PARAMETERS:
             n, the number of sites
             m, the number of species
             phi, the matrix fill, i.e. the proportion of nonzero elements
             T, the desired temperature of the matrix
             err, half-width of acceptable interval around T
                  e.g. err = 1 (default) stops function execution
                  when matrix temperature is within 1 degree of T
             max_iter, maximum iterations before function terminates
  RETURNS:
         A tuple containing the matrix and its actual temperature.
  NOTES:
         Some combinations of T and phi cause the algorithm
         to converge on a lower temperature than desired.
  """
  # First, create a perfectly nested matrix with dimensions n,m and fill phi
  mat = perfnest(n,m,phi)
  # keep track of successive versions of mat.  If the algorithm overshoots temperature T,
  # it will step back to a previous version of the matrix and try again.
  mat_versions = [mat.copy()]

  p = get_parameter(phi,n,m) # get isocline parameter
  temp = 0  # current temperature of mat
  prev_temps = [temp] # keep track of previous temperatures, to check for convergence
  # get unexpectedness associated with temperature T
  # T = k*U, where k = 100/Umax = 100/0.04145
  target_unexp = T / (100/0.04145)

  # counter and stopping condition
  i = 0
  done = False

  # Iteratively perturb and pack the matrix until either T or max_iter is reached
  while done is False and i < max_iter:  # loop level 1
    i = i + 1
    #print i
    # calculate current unexpectedness
    unexp = totUnexp(mat, n, m, p)
    cols_done = False # columns contain cells to be flipped

    # perturb matrix by flipping cells until either target_unexp
    # is reached or no more cells are left to flip.
    j = 0
    while unexp < target_unexp and cols_done is False: # loop level 2
      j = j + 1

      col_indices = range(m) # all possible columns
      need_new_index = True
      # Randomly select columns to perturb until either a successful flip occurs
      # or no more columns are left to try
      while need_new_index is True:  # loop level 3
        col_index = random.sample(col_indices,1)[0]
        col_flipped = flipCol(mat, col_index, n, m, p)

        if col_flipped != -1: # flip was successful
          need_new_index = False
        else: # no flip; one side of isocline has been completely flipped already
          if len(col_indices) > 1:
            col_indices.remove(col_index) # remove col_index from sample pool
          else:
            col_indices.remove(col_index)
            need_new_index = False        # no more columns left to choose
      # end loop level 3

      # If all columns have been tried unsuccessfully, break out of loop level 2,
      # i.e. go directly to matrix packing and temperature calculation.
      if len(col_indices) == 0:
        cols_done = True
        continue
      else: # flip cells in rows that were affected by the column flip; get unexpectedness
        row_flipped = flipRow(mat, col_flipped[0][0], n, m, p)
        row_flipped = flipRow(mat, col_flipped[1][0], n, m, p)

        # get total unexpectedness = sum of cells in the unexpectedness matrix
        unexp = totUnexp(mat,n,m,p)

        if unexp > target_unexp:  # more than enough unexpectedness. break out of loop level 2
          cols_done = True
          continue
        # else: keep perturbing the matrix
    # end loop level 2

    # Now the matrix has either reached target_unexp or cannot be further perturbed.
    # pack the matrix and get its temperature.
    temp_info = temperature(mat)
    temp = temp_info[0]
    mat  = temp_info[1]

    # Check for early convergence and stopping conditions
    if temp == prev_temps[-1]: # matrix has converged on a temperature less than T
      # step back to the previous matrix
      mat = mat_versions[-1].copy()
      # discard most recent matrix unless only the first matrix remains
      if len(mat_versions) > 1:
        mat_versions = mat_versions[:-1]
      #prev_temps   = prev_temps[:-1]
      # don't discard most recent temperature- otherwise the algorithm won't recognize
      # early convergence
    elif temp < T-err: # matrix temperature is below T
      # update mat_versions and prev_temps
      mat_versions.append(mat.copy())
      prev_temps.append(temp)
    elif temp > (T-err) and temp < (T+err):  # matrix temp is close enough to T. terminate loop.
      done = True
    else: # temp > T+err; matrix temperature is too high.
      # step back to previous matrix (unless the current matrix is the first matrix)
      if len(mat_versions) > 1 and len(prev_temps) > 1:
        mat = mat_versions[-1].copy()
        mat_versions = mat_versions[:-1]
        prev_temps   = prev_temps[:-1]

    if i == max_iter:  # user-defined maximum iterations have been reached
        # return the matrix with the highest unexpectedness
        # print "choosing best matrix..."
        unexps = []
        if len(mat_versions) > 1:
          for mv in mat_versions:
            unexps.append(totUnexp(mv,n,m,p))
          best_unexp = max(unexps)
          best_index = unexps.index(best_unexp)
          mat = mat_versions[best_index]
          temp =  temperature(mat)[0]
        # else, the only previous matrix to choose from is the perfectly nested one.
        # stick with the most recent matrix.

  # end loop level 1

  return (mat, temp)
# end hotMess()